"""
数据库迁移脚本：为用户表添加班级和真实姓名字段
执行时间：2024年
"""

import sqlite3
import os
import sys

def upgrade():
    """
    升级数据库：添加班级和真实姓名字段
    """
    print("开始数据库迁移：添加班级和真实姓名字段...")
    
    # 数据库文件路径
    db_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'instance', 'mailserver_dev.db')
    
    if not os.path.exists(db_path):
        print(f"❌ 数据库文件不存在: {db_path}")
        return
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查字段是否已存在
        cursor.execute("PRAGMA table_info(users)")
        columns = [col[1] for col in cursor.fetchall()]
        
        # 添加real_name字段
        if 'real_name' not in columns:
            cursor.execute('ALTER TABLE users ADD COLUMN real_name VARCHAR(50)')
            print("✓ 已添加 real_name 字段")
        else:
            print("✓ real_name 字段已存在")
        
        # 添加class_name字段
        if 'class_name' not in columns:
            cursor.execute('ALTER TABLE users ADD COLUMN class_name VARCHAR(50)')
            print("✓ 已添加 class_name 字段")
        else:
            print("✓ class_name 字段已存在")
        
        # 为class_name字段创建索引
        try:
            cursor.execute('CREATE INDEX IF NOT EXISTS ix_users_class_name ON users (class_name)')
            print("✓ 已为 class_name 字段创建索引")
        except Exception as e:
            print(f"⚠ 创建索引时出现警告: {e}")
        
        conn.commit()
        conn.close()
        
        print("数据库迁移完成！")
        
    except Exception as e:
        print(f"❌ 数据库迁移失败: {e}")
        if 'conn' in locals():
            conn.close()
        raise

def downgrade():
    """
    降级数据库：移除班级和真实姓名字段
    注意：SQLite不支持DROP COLUMN，需要重建表
    """
    print("开始数据库回滚：移除班级和真实姓名字段...")
    
    # 数据库文件路径
    db_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'instance', 'mailserver_dev.db')
    
    if not os.path.exists(db_path):
        print(f"❌ 数据库文件不存在: {db_path}")
        return
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 创建临时表（不包含新字段）
        cursor.execute('''
            CREATE TABLE users_temp (
                id INTEGER PRIMARY KEY,
                username VARCHAR(80) UNIQUE NOT NULL,
                email VARCHAR(120) UNIQUE NOT NULL,
                password_hash VARCHAR(255) NOT NULL,
                role VARCHAR(20) NOT NULL DEFAULT 'student',
                is_active BOOLEAN NOT NULL DEFAULT 1,
                created_at DATETIME NOT NULL,
                last_login DATETIME
            )
        ''')
        
        # 复制数据到临时表
        cursor.execute('''
            INSERT INTO users_temp (id, username, email, password_hash, role, is_active, created_at, last_login)
            SELECT id, username, email, password_hash, role, is_active, created_at, last_login
            FROM users
        ''')
        
        # 删除原表
        cursor.execute('DROP TABLE users')
        
        # 重命名临时表
        cursor.execute('ALTER TABLE users_temp RENAME TO users')
        
        # 重建索引
        cursor.execute('CREATE UNIQUE INDEX ix_users_username ON users (username)')
        cursor.execute('CREATE UNIQUE INDEX ix_users_email ON users (email)')
        
        conn.commit()
        conn.close()
        
        print("✓ 已移除 real_name 和 class_name 字段")
        print("数据库回滚完成！")
        
    except Exception as e:
        print(f"❌ 数据库回滚失败: {e}")
        if 'conn' in locals():
            conn.close()
        raise

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='数据库迁移脚本')
    parser.add_argument('action', choices=['upgrade', 'downgrade'], 
                       help='迁移操作：upgrade(升级) 或 downgrade(降级)')
    
    args = parser.parse_args()
    
    if args.action == 'upgrade':
        upgrade()
    elif args.action == 'downgrade':
        downgrade()